#include "non_threshold_evaluator.h"

#include <fstream>
#include <vector>
#include <string>
#include <iostream>
#include <strstream>
#include <algorithm>
using namespace std;

void read_label(string label_file, vector<int> & label_list) {
	ifstream fin(label_file.c_str());
	if(fin.fail()) {
		cout << "Cannnot open " << label_file << endl;
		return;
	}

	while(!fin.eof()) {
		string line;
		getline(fin, line);
		if(line.empty()) {
			break;
		}
		
		int label = line[0] - '0';
		
		if(label != 1) {
			label = 0;
		}
		label_list.push_back(label);
	}
	fin.close();
}

void read_score(string score_file, vector<double> & score_list) {
	ifstream fin(score_file.c_str());
	if(fin.fail()) {
		cout << "Cannnot open " << score_file << endl;
		return;
	}

	string line;
	getline(fin, line);
	while(!fin.eof()) {
		getline(fin, line);
		if(line.empty()) {
			break;
		}

		int start = 2;
		int end = line.find(' ', start);

		double score = atof(line.substr(start, end - start).c_str());

		score_list.push_back(score);
	}
	fin.close();
}

void read_score_batch(string score_file, vector<vector<double> > & score_list) {
	ifstream fin(score_file.c_str());
	if(fin.fail()) {
		cout << "Cannnot open " << score_file << endl;
		return;
	}

	string line;
	while(!fin.eof()) {
		getline(fin, line);
		if(line.empty()) {
			break;
		}

		int start = 0;
		int end = 0;
		for(int i = 0; end < (int)line.size(); i++, start = end + 1) {
			end = line.find('\t', start);
			if(end == -1) {
				end = (int)line.size();
			}

			double score = atof(line.substr(start, end - start).c_str());

			if((unsigned)i >= score_list.size()) {
				vector<double> l;
				l.push_back(score);
				score_list.push_back(l);
			} else {
				score_list[i].push_back(score);
			}
		}
	}
	fin.close();
}

struct pair_comp {
	bool operator() (pair<int, double> * i, pair<int, double> * j) { return (i->second > j->second);}
};

void sort_by_score(vector<int> & label_list, vector<double> & score_list) {
	vector<pair<int, double> *> pair_list;
	for(unsigned i = 0; i < label_list.size(); i++) {
		pair<int, double> * p = new pair<int, double>();
		p->first = label_list[i];
		p->second = score_list[i];
		pair_list.push_back(p);
	}

	random_shuffle(pair_list.begin(), pair_list.end());

	pair_comp pc;
	sort (pair_list.begin(), pair_list.end(), pc);

	for(unsigned i = 0; i < pair_list.size(); i++) {
		label_list[i] = pair_list[i]->first;
		score_list[i] = pair_list[i]->second;
	}

	for(unsigned i = 0; i < pair_list.size(); i++) {
		delete pair_list[i];
	}
}

void proceed_batch(string input, string label, string output, string type, int bin) {
	vector<int> label_list;
	read_label(label, label_list);

	vector<vector<double> > score_list;
	read_score_batch(input, score_list);

	if(label_list.size() != score_list[0].size()) {
		cout << "Length mismatch!" << endl;
		return;
	}

	ofstream fout((output + ".txt").c_str());
	non_threshold_evaluator * evaluator = get_non_threshold_evaluator(type);
	for(unsigned i = 0; i < score_list.size(); i++) {
		vector<double> scores = score_list[i];
		vector<int> labels = label_list;
		sort_by_score(labels, scores);
		
		strstream sstr;
		sstr << output << "." << i << ".txt";
		string output_layer;
		getline(sstr, output_layer);

		double r = evaluator->evaluate(labels, scores, output_layer, bin);
		fout << r << endl;
	}
	fout.close();
	delete evaluator;
}

void proceed(string input, string label, string output, string type, int bin) {
	vector<int> label_list;
	vector<double> score_list;

	read_label(label, label_list);
	read_score(input, score_list);

	if(label_list.size() != score_list.size()) {
		cout << "Length mismatch!" << endl;
		return;
	}

	sort_by_score(label_list, score_list);

	non_threshold_evaluator * evaluator = get_non_threshold_evaluator(type);
	evaluator->evaluate(label_list, score_list, output, bin);
	delete evaluator;
}

int main(int argc, char ** argv) {
	if(argc != 7) {
		return 0;
	}

	string input = argv[1];
	string label = argv[2];
	string output = argv[3];
	string type = argv[4];
	int bin = atoi(argv[5]);
	bool batch_mode = atoi(argv[6]) == 1;

	if(batch_mode) {
		proceed_batch(input, label, output, type, bin);
	} else {
		proceed(input, label, output, type, bin);
	}

	return 1;
}
